import numpy as np
import time
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import argparse
# Import seaborn
#import seaborn as sns

parser = argparse.ArgumentParser(description='Thermal evolution model of the EL planetesimal with breakup')
parser.add_argument("-b","--breakuptime",type=int,help='Time of the EL parent breakup',default=180)
parser.add_argument("-t","--totalSimualtionTime",type=float,help='Total simulation time',default=200)
parser.add_argument("-p", "--plot",action="store_true",help="Show plots")
parser.add_argument("-s", "--save",action="store_true",help="Save temperature curves")
parser.add_argument("-o","--outsidetemperature",type=int,help='Temperature at surface',default=200)
parser.add_argument("-n","--iterations",type=float,help='Number of Monte Carlo iterations',default=1000000)
parser.add_argument("-0", "--best",action="store_true",help="Calc Χ^2 at best")

istacoolcolor='orchid'
args = parser.parse_args()

# Setup a plot such that only the bottom spine is shown
def setup(ax):
    #ax.spines['right'].set_color('none')
    #ax.spines['left'].set_color('none')
    #ax.yaxis.set_major_locator(ticker.NullLocator())
    #ax.spines['top'].set_color('none')
    #ax.xaxis.set_ticks_position('bottom')
    ax.tick_params(which='major', width=1.00)
    ax.tick_params(which='major', length=12)
    ax.tick_params(which='minor', width=0.75)
    ax.tick_params(which='minor', length=2.5)
    #ax.set_xlim(0, 5)
    #ax.set_ylim(0, 1)
    #ax.patch.set_alpha(0.0)

breakupTime=200
fname='eltemp-'+str("{:03d}".format(int(breakupTime)))+'.dat'
tV,TV = np.loadtxt(fname, unpack=True)
tempRef=TV
timeRef=tV

dtype1 = np.dtype([('mateorite', '|S20'), ('thermochronometer', '|S7'), \
                       ('age', 'f8'), ('ageUncertainty','f8'), \
                       ('temperature', 'f8'), ('temperatureUncertainty','f8')])
metName, tcName, mAge, sAge, mTemp, msTemp = np.loadtxt('tagesV3.dat', \
                            usecols=(0,1,2,4,5,7), dtype=dtype1,unpack=True)


plt.xlabel('Time after CAIs (Myr)')
plt.ylabel('Temperature (K)')

breakupTime=200
fname='eltemp-'+str("{:03d}".format(int(breakupTime)))+'.dat'
tV,TV = np.loadtxt(fname, unpack=True)
modelLineNoBreak, = plt.plot(tV,TV,'-',c='#000000', zorder=3)

######### BREAK UP ###########################################################
## color palette from https://davidmathlogic.com/colorblind/#%23648FFF-%23785EF0-%23DC267F-%23FE6100-%23FFB000
## click on the IBM Design Library... this is to help color blind people 
breakupTime=90
tV=np.arange(5,args.totalSimualtionTime)
#### plot insta caoling
TV=np.interp(tV, timeRef, tempRef)
TV[np.where(tV>breakupTime)]=args.outsidetemperature
modelLine90MyrReaccumulated, = plt.plot(tV,TV,':',c='fuchsia',zorder=-1)

fname='eltemp-'+str("{:03d}".format(int(breakupTime)))+'.dat'
tV,TV = np.loadtxt(fname, unpack=True)
modelLine90MyrMonolithic, = plt.plot(tV,TV,'--',c='fuchsia')

breakupTime=59
#fname='istacool-'+str("{:03d}".format(int(breakupTime)))+'.dat'
#tV,TV = np.loadtxt(fname, unpack=True)
#plt.plot(tV,TV,'-.',c=istacoolcolor,zorder=-1)
#### plot insta caoling
TV=np.interp(tV, timeRef, tempRef)
TV[np.where(tV>breakupTime)]=args.outsidetemperature
modelLine59MyrReaccumulated, = plt.plot(tV,TV,':',c='darkviolet',zorder=-1)
#### plot slower cooling
breakupTime=57
fname='eltemp-'+str("{:03d}".format(int(breakupTime)))+'.dat'
tV,TV = np.loadtxt(fname, unpack=True)
modelLine59MyrMonolithic, = plt.plot(tV,TV,'--',c='darkturquoise')


'''
breakupTime=48
fname='eltemp-'+str("{:03d}".format(int(breakupTime)))+'.dat'
tV,TV = np.loadtxt(fname, unpack=True)
plt.plot(tV,TV,c='k')
'''

breakupTime=40
#fname='istacool-'+str("{:03d}".format(int(breakupTime)))+'.dat'
#tV,TV = np.loadtxt(fname, unpack=True)
#plt.plot(tV,TV,c=istacoolcolor,zorder=-1)
#### plot insta caoling
TV=np.interp(tV, timeRef, tempRef)
TV[np.where(tV>breakupTime)]=args.outsidetemperature
modelLine40MyrReaccumulated, = plt.plot(tV,TV,':',c='blue',zorder=-1)

fname='eltemp-'+str("{:03d}".format(int(breakupTime)))+'.dat'
tV,TV = np.loadtxt(fname, unpack=True)
modelLine40MyrMonolithic, = plt.plot(tV,TV,'--',c='blue')

'''
breakupTime=10
#TV=np.interp(tV, timeRef, tempRef)
#TV[np.where(tV>breakupTime)]=args.outsidetemperature
#plt.plot(tV,TV,c=istacoolcolor,zorder=-1)

fname='eltemp-'+str("{:03d}".format(int(breakupTime)))+'.dat'
tV,TV = np.loadtxt(fname, unpack=True)
plt.plot(tV,TV,c='k')
'''
'''
for k,s,c in zip([b'I-Xe', b'Ar-Ar', b'Rb-Sr'], ['o','^','s'], ['coral','blue','green']):
    j=np.argwhere(tcName == k)
    j=j.reshape(len(j))
    plt.scatter(4567.-mAge[j], mTemp[j], marker=s, c='k', size=7, zorder=4)
'''
for k,s,c in zip([b'I-Xe', b'Ar-Ar', b'Rb-Sr'], ['o','^','s'],
                     ['#F30052','#FFA000','#00a73f']): ### colors of the data points
    j=np.argwhere(tcName == k)
    j=j.reshape(len(j))
    plt.errorbar(4567.-mAge[j], mTemp[j], fmt=s, capsize=3, ms=7, mew=0.5, \
                     xerr=sAge[j], yerr=msTemp[j], c=c, mec='k',
                     label=k.decode(), zorder=25)


legend1 = plt.legend([modelLineNoBreak,
                          modelLine90MyrReaccumulated, modelLine90MyrMonolithic, 
                          modelLine59MyrReaccumulated, modelLine59MyrMonolithic,
                          modelLine40MyrReaccumulated, modelLine40MyrMonolithic],
                            ["Breakup > 200 Myr",
                            "Breakup @ 90Myr, reaccumulated",
                            "Breakup @ 90Myr, monolithic",
                            "Breakup @ 59Myr, reaccumulated",
                            "Breakup @ 57Myr, monolithic",
                            "Breakup @ 40Myr, reaccumulated",
                            "Breakup @ 40Myr, monolithic"],
                         loc=0, fontsize=10)
plt.gca().add_artist(legend1)

plt.legend(loc=3, fontsize=12)
plt.ylim((190,1110))
plt.xlim((7,200))
plt.xscale('log')
#plt.yscale('log')
plt.tight_layout()
ax=plt.gca()
setup(ax)
#subs = [7, 8, 9, 20, 30, 50, 70]  # ticks to show per decade
ax.yaxis.set_major_locator(ticker.FixedLocator(np.arange(200,1200,100))) #set the ticks position
ax.xaxis.set_major_formatter(ticker.ScalarFormatter())
ax.xaxis.set_major_formatter(ticker.ScalarFormatter())
ax.xaxis.set_minor_formatter(ticker.ScalarFormatter())
#ax.xaxis.get_major_formatter().set_scientific(False)
#ax.xaxis.get_major_formatter().set_useOffset(False)
#ax.xaxis.set_minor_locator(ticker.FixedLocator(subs)) #set the ticks position
#ax.ticklabel_format(style='plain', axis='x')
#ax.ticklabel_format(axis="x", style="plain", useOffset=True, scilimits=(0,0))
#ax.xaxis.set_label_coords(0, 0)
ax.text(45, 310,  "0%", c='blue', rotation=-80, size=12)
ax.text(60, 310,  "5%", c='darkturquoise', rotation=-70, size=12)#size="x-small",
ax.text(93, 270, "68%", c='fuchsia', rotation=-80, size=12) ## yellow at 90 Myr

ax.text(36, 220,  "0%", c='blue', rotation=90, size=12)
ax.text(53, 220,  "5%", c='darkviolet', rotation=90, size=12)
ax.text(82, 220,  "68%", c='fuchsia', rotation=90, size=12) ## yellow at 90 Myr

#ax.text(40, 580,  u'$\leftarrow$40Myr', rotation=45, size=12)
#ax.text(60, 480,  u'$\leftarrow$60Myr', rotation=45, size=12)
#ax.text(93, 367,  u'$\leftarrow$90Myr', rotation=0, size=12)

plt.savefig('figS1.pdf', dpi=300)
#plt.show()
